"""
Batch Processing Example for Wan 2.2 API

This script demonstrates how to submit multiple video generation tasks in parallel
and handle the results efficiently.
"""
import asyncio
import os
import random
from typing import List, Dict, Any
from datetime import datetime
from pathlib import Path

# Add the parent directory to the path so we can import the client
import sys
sys.path.insert(0, str(Path(__file__).parent.parent))

from client import WanClient, TaskStatus, T2VRequest

# Configuration
API_BASE_URL = "http://localhost:8000"  # Update with your API URL
API_KEY = os.getenv("WAN_API_KEY")      # Optional, if authentication is enabled
CONCURRENT_TASKS = 3                     # Number of tasks to run in parallel

# Sample prompts for batch processing
SAMPLE_PROMPTS = [
    "A futuristic cityscape at night, neon lights, cyberpunk style, 8K",
    "A serene mountain lake at sunrise, misty atmosphere, nature photography",
    "An astronaut floating in space with Earth in the background, highly detailed",
    "A steampunk airship flying through clouds, intricate details, cinematic",
    "A magical forest with glowing plants and fireflies, fantasy style"
]

class BatchProcessor:
    def __init__(self, client: WanClient, max_concurrent: int = 3):
        self.client = client
        self.max_concurrent = max_concurrent
        self.semaphore = asyncio.Semaphore(max_concurrent)
        self.completed_tasks: List[Dict[str, Any]] = []
        self.failed_tasks: List[Dict[str, Any]] = []
    
    async def process_prompt(self, prompt: str, index: int) -> None:
        """Process a single prompt and track the result."""
        task_info = {
            "prompt": prompt,
            "index": index,
            "start_time": datetime.now(),
            "end_time": None,
            "status": "pending",
            "task_id": None,
            "result": None,
            "error": None
        }
        
        try:
            async with self.semaphore:
                print(f"[{index}] Submitting task: {prompt[:50]}...")
                task = await self.client.text_to_video(
                    prompt=prompt,
                    duration_seconds=random.randint(3, 6),  # Random duration between 3-6 seconds
                    resolution=random.choice(["720p", "1080p"]),
                    seed=random.randint(0, 10000),
                    wait=False
                )
                task_info["task_id"] = task.task_id
                task_info["status"] = "submitted"
                
                # Wait for task completion
                try:
                    result = await self.client.wait_for_task_completion(
                        task.task_id,
                        poll_interval=2.0,
                        timeout=300  # 5 minutes timeout per task
                    )
                    task_info["end_time"] = datetime.now()
                    
                    if result.status == TaskStatus.COMPLETED:
                        task_info["status"] = "completed"
                        task_info["result"] = result.result
                        self.completed_tasks.append(task_info)
                        print(f"✅ [{index}] Task completed successfully!")
                    else:
                        task_info["status"] = "failed"
                        task_info["error"] = result.error
                        self.failed_tasks.append(task_info)
                        print(f"❌ [{index}] Task failed: {result.error}")
                        
                except asyncio.TimeoutError:
                    task_info["status"] = "timeout"
                    task_info["error"] = "Task timed out"
                    self.failed_tasks.append(task_info)
                    print(f"⏱️  [{index}] Task timed out")
        
        except Exception as e:
            task_info["status"] = "error"
            task_info["error"] = str(e)
            self.failed_tasks.append(task_info)
            print(f"⚠️  [{index}] Error processing task: {str(e)}")
    
    async def process_batch(self, prompts: List[str]) -> Dict[str, Any]:
        """Process a batch of prompts in parallel."""
        start_time = datetime.now()
        print(f"Starting batch processing of {len(prompts)} tasks with {self.max_concurrent} concurrent tasks...")
        
        # Create and gather all tasks
        tasks = [
            self.process_prompt(prompt, i)
            for i, prompt in enumerate(prompts, 1)
        ]
        await asyncio.gather(*tasks)
        
        # Calculate statistics
        end_time = datetime.now()
        total_time = (end_time - start_time).total_seconds()
        
        return {
            "total_tasks": len(prompts),
            "completed": len(self.completed_tasks),
            "failed": len(self.failed_tasks),
            "success_rate": len(self.completed_tasks) / len(prompts) * 100 if prompts else 0,
            "total_time_seconds": total_time,
            "tasks_per_minute": len(prompts) / (total_time / 60) if total_time > 0 else 0,
            "start_time": start_time.isoformat(),
            "end_time": end_time.isoformat()
        }

async def main():
    # Initialize the client
    print("Initializing Wan 2.2 API client...")
    client = WanClient(base_url=API_BASE_URL, api_key=API_KEY)
    
    # Create output directory
    output_dir = Path("batch_output")
    output_dir.mkdir(exist_ok=True)
    
    try:
        # Initialize batch processor
        processor = BatchProcessor(client, max_concurrent=CONCURRENT_TASKS)
        
        # Process the batch
        print(f"\n=== Starting Batch Processing ===")
        stats = await processor.process_batch(SAMPLE_PROMPTS)
        
        # Print summary
        print("\n=== Batch Processing Complete ===")
        print(f"Total tasks: {stats['total_tasks']}")
        print(f"Completed: {stats['completed']}")
        print(f"Failed: {stats['failed']}")
        print(f"Success rate: {stats['success_rate']:.1f}%")
        print(f"Total time: {stats['total_time_seconds']:.1f} seconds")
        print(f"Tasks per minute: {stats['tasks_per_minute']:.1f}")
        
        # Save results to a file
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        results_file = output_dir / f"batch_results_{timestamp}.txt"
        
        with open(results_file, "w", encoding="utf-8") as f:
            f.write("=== Batch Processing Results ===\n\n")
            f.write(f"Start time: {stats['start_time']}\n")
            f.write(f"End time: {stats['end_time']}\n")
            f.write(f"Total tasks: {stats['total_tasks']}\n")
            f.write(f"Completed: {stats['completed']}\n")
            f.write(f"Failed: {stats['failed']}\n")
            f.write(f"Success rate: {stats['success_rate']:.1f}%\n")
            f.write(f"Total time: {stats['total_time_seconds']:.1f} seconds\n")
            f.write(f"Tasks per minute: {stats['tasks_per_minute']:.1f}\n\n")
            
            f.write("\n=== Completed Tasks ===\n")
            for task in processor.completed_tasks:
                duration = (task['end_time'] - task['start_time']).total_seconds()
                f.write(f"\n[{task['index']}] {task['prompt'][:100]}...\n")
                f.write(f"    Status: {task['status'].upper()}\n")
                f.write(f"    Task ID: {task['task_id']}\n")
                f.write(f"    Duration: {duration:.1f} seconds\n")
                if task['result'] and 'video_url' in task['result']:
                    f.write(f"    Video URL: {task['result']['video_url']}\n")
            
            if processor.failed_tasks:
                f.write("\n=== Failed Tasks ===\n")
                for task in processor.failed_tasks:
                    f.write(f"\n[{task['index']}] {task['prompt'][:100]}...\n")
                    f.write(f"    Status: {task['status'].upper()}\n")
                    if task.get('task_id'):
                        f.write(f"    Task ID: {task['task_id']}\n")
                    if task['error']:
                        f.write(f"    Error: {task['error']}\n")
        
        print(f"\nDetailed results saved to: {results_file}")
    
    except Exception as e:
        print(f"An error occurred: {str(e)}")
        import traceback
        traceback.print_exc()
    finally:
        await client.close()

if __name__ == "__main__":
    asyncio.run(main())
